// Copyright Kani Contributors
// SPDX-License-Identifier: Apache-2.0 OR MIT

use crate::CbmcSolver;
use serde::{Deserialize, Serialize};
use std::{borrow::Borrow, collections::BTreeSet, path::PathBuf};
use strum_macros::Display;
use tracing::{debug, trace};

/// A CBMC-level `assigns` contract that needs to be enforced on a function.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct AssignsContract {
    /// The target of the contract
    pub contracted_function_name: String,
    /// A static global variable used to track recursion that must not be havocked.
    /// This is only needed if the function is tagged with `#[kani::recursive]`
    pub recursion_tracker: Option<String>,
}

/// We emit this structure for each annotated proof harness (`#[kani::proof]`) we find.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct HarnessMetadata {
    /// The fully qualified name the user gave to the function (i.e. includes the module path).
    pub pretty_name: String,
    /// The name of the function in the CBMC symbol table.
    pub mangled_name: String,
    /// The name of the crate this harness belongs to.
    pub crate_name: String,
    /// The (currently full-) path to the file this proof harness was declared within.
    pub original_file: String,
    /// The line in that file where the proof harness begins.
    pub original_start_line: usize,
    /// The line in that file where the proof harness ends.
    pub original_end_line: usize,
    /// Optional modeling file that was generated by the compiler that includes this harness.
    pub goto_file: Option<PathBuf>,
    /// The `#[kani::<>]` attributes added to a harness.
    pub attributes: HarnessAttributes,
    /// A CBMC-level assigns contract that should be enforced when running this harness.
    pub contract: Option<AssignsContract>,
    /// If the harness contains some usage of loop contracts.
    pub has_loop_contracts: bool,
    /// If the harness was automatically generated or manually written.
    pub is_automatically_generated: bool,
}

/// The attributes added by the user to control how a harness is executed.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct HarnessAttributes {
    /// Whether the harness has been annotated with proof.
    pub kind: HarnessKind,
    /// Whether the harness is expected to panic or not.
    pub should_panic: bool,
    /// Optional data to store solver.
    pub solver: Option<CbmcSolver>,
    /// Optional data to store unwind value.
    pub unwind_value: Option<u32>,
    /// The stubs used in this harness.
    pub stubs: Vec<Stub>,
    /// The name of the functions being stubbed by their contract.
    pub verified_stubs: Vec<String>,
}

#[derive(Clone, Eq, PartialEq, Debug, Display, Serialize, Deserialize)]
pub enum HarnessKind {
    /// Function was annotated with `#[kani::proof]`.
    #[strum(serialize = "#[kani::proof]")]
    Proof,
    /// Function was annotated with `#[kani::proof_for_contract(target_fn)]`.
    #[strum(serialize = "#[kani::proof_for_contract]")]
    ProofForContract { target_fn: String },
    /// This is a test harness annotated with `#[test]`.
    #[strum(serialize = "#[test]")]
    Test,
}

impl HarnessAttributes {
    /// Create a new harness of the provided kind.
    pub fn new(kind: HarnessKind) -> HarnessAttributes {
        HarnessAttributes {
            kind,
            should_panic: false,
            solver: None,
            unwind_value: None,
            stubs: vec![],
            verified_stubs: vec![],
        }
    }

    /// Return whether this is a proof harness.
    pub fn is_proof_harness(&self) -> bool {
        matches!(self.kind, HarnessKind::Proof | HarnessKind::ProofForContract { .. })
    }
}

/// The stubbing type.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Stub {
    pub original: String,
    pub replacement: String,
}

impl HarnessMetadata {
    /// get the unqualifed (i.e. without ::) harness name. If the
    /// harness name contains ::, then we use rightmost name..
    pub fn get_harness_name_unqualified(&self) -> &str {
        const PATH_SEPARATOR: &str = "::";
        if let Some(last_separator) = self.pretty_name.rfind(PATH_SEPARATOR) {
            let name_start = last_separator + PATH_SEPARATOR.len();
            &self.pretty_name[name_start..]
        } else {
            &self.pretty_name
        }
    }
}

/// Search for a proof harness with a particular name.
/// At the present time, we use `no_mangle` so collisions shouldn't happen,
/// but this function is written to be robust against that changing in the future.
pub fn find_proof_harnesses<'a, I>(
    targets: &BTreeSet<&String>,
    all_harnesses: I,
    exact_filter: bool,
) -> Vec<&'a HarnessMetadata>
where
    I: IntoIterator,
    I::Item: Borrow<&'a HarnessMetadata>,
{
    debug!(?targets, "find_proof_harness");
    let mut result = vec![];
    for md in all_harnesses.into_iter() {
        let md: &'a HarnessMetadata = md.borrow();

        // --harnesses should not select automatic harnesses
        if md.is_automatically_generated {
            continue;
        }
        if exact_filter {
            // Check for exact match only
            if targets.contains(&md.pretty_name) {
                // if exact match found, stop searching
                result.push(md);
            } else {
                trace!(skip = md.pretty_name, "find_proof_harnesses");
            }
        } else {
            // Either an exact match, or a substring match. We check the exact first since it's cheaper.
            if targets.contains(&md.pretty_name)
                || targets.contains(&md.get_harness_name_unqualified().to_string())
                || targets.iter().any(|target| md.pretty_name.contains(*target))
            {
                result.push(md);
            } else {
                trace!(skip = md.pretty_name, "find_proof_harnesses");
            }
        }
    }
    result
}
